# -*- coding: utf-8 -*-
# @Author  : gaoyu
# @Time    : 2024/9/26
# @Function: SM4加密工具函数

import base64
import json
import os
import sys
from typing import Dict, Any

import pandas as pd
from gmssl import sm4
from loguru import logger

# 配置loguru日志格式，使其更加花里胡哨
logger.remove()
logger.add(
    sys.stderr,
    format="<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | "
           "<level>{level: <8}</level> | "
           "<cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> | "
           "<level>{message}</level>",
    level="INFO",
    colorize=True,
)

from hydropostgis.utils.sm2.pairwise_crypto import PairwiseCrypto, log_header, log_subheader


def encrypt_with_sm4(data: str, key: bytes) -> bytes:
    """
    使用SM4算法加密数据

    Args:
        data: 要加密的数据
        key: SM4密钥

    Returns:
        bytes: 加密后的数据
    """
    # 使用SM4加密数据
    sm4_crypt = sm4.CryptSM4()
    sm4_crypt.set_key(key, sm4.SM4_ENCRYPT)

    # 确保数据长度是16的倍数（SM4块大小），添加PKCS#7填充
    data_bytes = data.encode('utf-8')
    padding_length = 16 - (len(data_bytes) % 16)
    if padding_length < 16:
        data_bytes += bytes([padding_length]) * padding_length

    # 使用SM4算法的ECB模式加密数据
    # ECB模式将明文分成固定大小的块，然后对每个块独立加密
    # 注意：在实际应用中，ECB模式不够安全，生产环境应考虑使用CBC/CTR等模式
    return sm4_crypt.crypt_ecb(data_bytes)


def decrypt_with_sm4(encrypted_data: bytes, key: bytes) -> str:
    """
    使用SM4算法解密数据

    Args:
        encrypted_data: 加密后的数据
        key: SM4密钥

    Returns:
        str: 解密后的数据
    """
    # 使用重构的密钥解密数据
    sm4_crypt = sm4.CryptSM4()
    sm4_crypt.set_key(key, sm4.SM4_DECRYPT)
    decrypted_data = sm4_crypt.crypt_ecb(encrypted_data)

    # 移除PKCS#7填充
    if decrypted_data:
        padding_length = decrypted_data[-1]
        if padding_length < 16 and all(b == padding_length for b in decrypted_data[-padding_length:]):
            decrypted_data = decrypted_data[:-padding_length]

    # 解码为字符串
    return decrypted_data.decode('utf-8')


def create_personal_encryption(crypto: PairwiseCrypto, data: str, encrypted_package: Dict[str, Any],
                               prefix: str = 'personal', key_file: str = None) -> Dict[str, Any]:
    """
    为个人创建加密数据和密钥分片

    Args:
        crypto: PairwiseCrypto实例
        data: 要加密的数据
        encrypted_package: 现有的加密包，将添加个人加密数据
        prefix: 密钥分片的前缀名，默认为'personal'
        key_file: 可选的密钥文件路径，如果提供则从文件读取密钥

    Returns:
        Dict: 更新后的加密包
    """
    # 尝试从文件读取密钥，如果文件不存在则生成新密钥
    personal_key = None
    key_name = f'{prefix}_key'

    if key_file and os.path.exists(key_file):
        try:
            with open(key_file, 'r') as f:
                key_data = json.load(f)
                if key_name in key_data:
                    personal_key = base64.b64decode(key_data[key_name])
                    logger.info(f"✅ 从文件 {key_file} 成功读取 {prefix} 密钥")
        except Exception as e:
            logger.warning(f"⚠️ 从文件读取密钥失败: {e}，将生成新密钥")

    # 如果未能从文件读取密钥，则生成新密钥
    if personal_key is None:
        personal_key = crypto.generate_sm4_key()
        # 如果提供了文件路径，保存新生成的密钥
        if key_file:
            try:
                key_data = {}
                if os.path.exists(key_file):
                    with open(key_file, 'r') as f:
                        key_data = json.load(f)
                key_data[key_name] = base64.b64encode(personal_key).decode('utf-8')
                with open(key_file, 'w') as f:
                    json.dump(key_data, f, indent=2)
                logger.info(f"✅ 新生成的 {prefix} 密钥已保存到文件 {key_file}")
            except Exception as e:
                logger.warning(f"⚠️ 保存密钥到文件失败: {e}")

    # 分割密钥为两个分片
    share1, share2 = crypto.split_key_for_pair(personal_key)

    # 使用个人密钥加密数据
    encrypted_data = encrypt_with_sm4(data, personal_key)

    # 将个人加密数据和密钥分片添加到加密包中
    updated_package = encrypted_package.copy()
    updated_package[f'{prefix}_encrypted_data'] = base64.b64encode(encrypted_data).decode('utf-8')
    updated_package[f'{prefix}_share1'] = base64.b64encode(share1).decode('utf-8')
    updated_package[f'{prefix}_share2'] = base64.b64encode(share2).decode('utf-8')

    return updated_package


def create_crypto(key_file):
    """
    创建并初始化一个PairwiseCrypto实例，包括设置密钥和添加参与者
    
    该函数完成以下任务:
    1. 创建PairwiseCrypto实例
    2. 尝试从文件读取主密钥，如果不存在则生成新密钥
    3. 设置SM4密钥
    4. 将新生成的密钥保存到文件（如果适用）
    5. 添加所有参与者到加密系统
    
    Args:
        key_file: 密钥文件路径，用于读取或保存主密钥
        
    Returns:
        PairwiseCrypto: 初始化完成的加密工具实例
    """
    # 创建一对一密钥分享工具
    crypto = PairwiseCrypto()
    log_header("密钥分享方案演示")
    logger.info(f"数据拥有者: {owner_name}")
    logger.info(f"参与者列表:")
    for user_id, name in participants.items():
        logger.info(f"  🧑‍💻 {user_id}: {name}")

    # 尝试从文件读取主密钥
    main_key = None
    if key_file and os.path.exists(key_file):
        try:
            with open(key_file, 'r') as f:
                key_data = json.load(f)
                if 'main_key' in key_data:
                    main_key = base64.b64decode(key_data['main_key'])
                    logger.info(f"✅ 从文件 {key_file} 成功读取主密钥")
        except Exception as e:
            logger.warning(f"⚠️ 从文件读取主密钥失败: {e}，将生成新密钥")
    crypto.set_sm4_key(main_key)
    # 如果提供了文件路径且成功生成了新密钥，保存主密钥
    if key_file and main_key is None:
        try:
            key_data = {}
            if os.path.exists(key_file):
                with open(key_file, 'r') as f:
                    key_data = json.load(f)
            key_data['main_key'] = base64.b64encode(crypto.sm4_key).decode('utf-8')
            with open(key_file, 'w') as f:
                json.dump(key_data, f, indent=2)
            logger.info(f"✅ 新生成的主密钥已保存到文件 {key_file}")
        except Exception as e:
            logger.warning(f"⚠️ 保存主密钥到文件失败: {e}")
    # 动态添加参与者
    for user_id, user_name in participants.items():
        crypto.add_participant(user_id)
        logger.info(f"✅ 已添加参与者: {user_id}({user_name})")
    return crypto


def get_encrypted_package(crypto: PairwiseCrypto, key_file: str, data_df: pd.DataFrame) -> Dict[str, Any]:
    """
    获取或创建加密包，优先从文件加载已有加密包，如不存在则重新创建
    
    该函数完成以下任务:
    1. 检查是否存在已加密的数据文件
    2. 如果存在，直接加载并返回加密包
    3. 如果不存在，则加密数据并创建新的加密包
    4. 打印密钥分片分发情况
    
    Args:
        crypto: PairwiseCrypto实例
        key_file: 密钥文件路径
        data_df: 要加密的数据DataFrame
        
    Returns:
        Dict: 包含加密数据和密钥分片的加密包
    """
    encrypted_package = None
    data_encrypted_file = f'data_encrypted_{key_file}'

    # 尝试从文件加载已加密的数据包
    if key_file and os.path.exists(key_file):
        try:
            with open(key_file, 'r') as f:
                key_data = json.load(f)
                # 检查是否包含必要的密钥分片信息
                if ('owner_share1' in key_data and 'owner_share2' in key_data and
                        'owner_encrypted_data' in key_data and 'encrypted_data' in key_data):
                    logger.info(f"✅ 从文件 {key_file} 成功加载加密包")

                    # 构建加密包
                    encrypted_package = {
                        'encrypted_data': key_data['encrypted_data'],
                        'owner_encrypted_data': key_data['owner_encrypted_data'],
                        'owner_share1': key_data['owner_share1'],
                        'owner_share2': key_data['owner_share2'],
                        'owner_shares': key_data.get('owner_shares', {}),
                        'participant_shares': key_data.get('participant_shares', {})
                    }

                    # 设置加密数据到crypto对象
                    crypto.set_encrypted_data(base64.b64decode(key_data['encrypted_data']))

                    # 恢复密钥分片
                    if 'owner_shares' in key_data:
                        crypto.owner_shares = key_data['owner_shares']
                    if 'participant_shares' in key_data:
                        crypto.participant_shares = key_data['participant_shares']
        except Exception as e:
            logger.warning(f"⚠️ 从文件加载加密包失败: {e}，将创建新的加密包")

    # 如果未能加载加密包，则创建新的
    if encrypted_package is None:
        # 将DataFrame转换为JSON字符串用于加密
        data_json = data_df.to_json(orient='records')
        logger.info(f"原始数据: DataFrame包含 {len(data_df)} 行, {len(data_df.columns)} 列")

        # 先加密数据，不分配密钥分片
        crypto.encrypt_data(data=data_json)
        logger.info("🔒 数据已加密，准备动态添加参与者")

        # 获取完整的加密包
        encrypted_package = crypto.get_encrypted_package()

        # 使用create_personal_encryption方法为拥有者创建个人加密数据
        # 使用owner作为前缀，而不是特定的名称
        encrypted_package = create_personal_encryption(crypto, data_json, encrypted_package, 'owner', key_file)

        # 保存加密包到文件
        if key_file:
            try:
                key_data = {}
                if os.path.exists(key_file):
                    with open(key_file, 'r') as f:
                        key_data = json.load(f)

                # 保存加密数据和拥有者密钥分片
                key_data['encrypted_data'] = encrypted_package['encrypted_data']
                key_data['owner_encrypted_data'] = encrypted_package['owner_encrypted_data']
                key_data['owner_share1'] = encrypted_package['owner_share1']
                key_data['owner_share2'] = encrypted_package['owner_share2']

                # 保存密钥分片
                key_data['owner_shares'] = encrypted_package['owner_shares']
                key_data['participant_shares'] = encrypted_package['participant_shares']

                with open(key_file, 'w') as f:
                    json.dump(key_data, f, indent=2)
                logger.info(f"✅ 加密包已保存到文件 {key_file}")
            except Exception as e:
                logger.warning(f"⚠️ 保存加密包到文件失败: {e}")

    logger.success(f"✅🔒 数据加密完成")

    # 打印密钥分片分发情况
    log_subheader("密钥分片分发情况")
    logger.info(f"{owner_name}持有的密钥分片:")
    for user_id, share in encrypted_package.get('owner_shares', {}).items():
        user_name = participants.get(user_id, user_id)
        logger.info(f"🔑 与{user_id}({user_name})共享的分片: {share[:20]}...")

    log_subheader("各参与者持有的密钥分片:")
    for user_id, share in encrypted_package.get('participant_shares', {}).items():
        user_name = participants.get(user_id, user_id)
        logger.info(f"🔑 {user_id}({user_name})的分片: {share[:20]}...")

    return encrypted_package


def demo_sm4_crypto(data_df: pd.DataFrame, owner_name: str, participants: Dict[str, str], key_file: str = None):
    """
    演示SM4加密工具和一对一密钥分享方案的使用
    
    Args:
        data_df: 要加密的数据DataFrame
        owner_name: 数据拥有者名称，如果为None则使用默认名称
        participants: 参与者字典，如果为None则使用默认参与者
        key_file: 可选的密钥文件路径，如果提供则尝试从文件读取密钥
    """

    try:
        # 创建一对一密钥分享工具
        crypto = create_crypto(key_file)

        # 获取加密包，优先从文件加载，如不存在则创建新的加密包
        # 这样可以避免重复加密数据，提高效率，同时保持密钥一致性
        # 如果加密数据存在，会加载
        encrypted_package = get_encrypted_package(crypto, key_file, data_df)

        # 第一部分：数据拥有者使用自己的密钥解密
        log_header(f"第一部分：{owner_name}独立解密")
        try:
            # 解码拥有者的密钥分片
            owner_share1_bytes = base64.b64decode(encrypted_package['owner_share1'])
            owner_share2_bytes = base64.b64decode(encrypted_package['owner_share2'])

            # 重构密钥
            reconstructed_key = crypto.reconstruct_key(owner_share1_bytes, owner_share2_bytes)

            # 使用重构的密钥解密数据
            owner_encrypted_data_bytes = base64.b64decode(encrypted_package['owner_encrypted_data'])
            decrypted_data = decrypt_with_sm4(owner_encrypted_data_bytes, reconstructed_key)

            # 将解密后的JSON字符串转换回DataFrame
            decrypted_df = pd.read_json(decrypted_data, orient='records')
            logger.success(
                f"🎯 {owner_name}独立解密结果: DataFrame包含 {len(decrypted_df)} 行, {len(decrypted_df.columns)} 列")
            logger.info(f"解密后的DataFrame前5行:\n{decrypted_df.head().to_string()}")
        except Exception as e:
            logger.error(f"💥 {owner_name}独立解密失败: {e}")

        # 第二部分：数据拥有者与其他参与者合作解密
        log_header(f"第二部分：{owner_name}与参与者合作解密")

        # 场景1-5：数据拥有者与各参与者合作解密
        for i, (user_id, user_name) in enumerate(participants.items(), 1):
            log_subheader(f"场景{i}：{owner_name}和{user_id}({user_name})合作解密")
            try:
                owner_share = encrypted_package['owner_shares'][user_id]
                participant_share = encrypted_package['participant_shares'][user_id]
                decrypted_data = crypto.decrypt_data(encrypted_package, owner_share, participant_share)

                # 将解密后的JSON字符串转换回DataFrame
                decrypted_df = pd.read_json(decrypted_data, orient='records')
                logger.success(f"🎯 解密结果: DataFrame包含 {len(decrypted_df)} 行, {len(decrypted_df.columns)} 列")
                logger.info(f"解密后的DataFrame前5行:\n{decrypted_df.head().to_string()}")
            except Exception as e:
                logger.error(f"💥 解密失败: {e}")

        # 第三部分：参与者之间尝试合作解密（预期失败）
        log_header("第三部分：参与者之间尝试合作解密")

        # 选择两个参与者进行测试
        user1_id, user2_id = list(participants.keys())[:2]
        user1_name = participants[user1_id]
        user2_name = participants[user2_id]

        # 场景6：第一个和第二个参与者尝试合作解密
        log_subheader(f"场景6：{user1_id}({user1_name})和{user2_id}({user2_name})尝试合作解密（应该失败）")
        try:
            owner_share = encrypted_package['participant_shares'][user1_id]
            participant_share = encrypted_package['participant_shares'][user2_id]
            decrypted_data = crypto.decrypt_data(encrypted_package, owner_share, participant_share)
            logger.warning(f"⚠️ 意外成功解密: {decrypted_data}")
        except Exception as e:
            logger.success(f"✅ 解密失败（符合预期）: {e}")

        # 选择另外两个参与者进行测试
        user3_id, user4_id = list(participants.keys())[2:4]
        user3_name = participants[user3_id]
        user4_name = participants[user4_id]

        # 场景7：第三个和第四个参与者尝试合作解密
        log_subheader(f"场景7：{user3_id}({user3_name})和{user4_id}({user4_name})尝试合作解密（应该失败）")
        try:
            owner_share = encrypted_package['participant_shares'][user3_id]
            participant_share = encrypted_package['participant_shares'][user4_id]
            decrypted_data = crypto.decrypt_data(encrypted_package, owner_share, participant_share)
            logger.warning(f"⚠️ 意外成功解密: {decrypted_data}")
        except Exception as e:
            logger.success(f"✅ 解密失败（符合预期）: {e}")

        log_header("演示完成")
        logger.success(f"🎉 所有场景测试完成")

    except Exception as e:
        logger.error(f"❌ 程序执行出错: {e}")


if __name__ == "__main__":
    # 创建示例DataFrame
    data = {
        "id": [1, 2, 3, 4, 5],
        "name": ["张三", "李四", "王五", "赵六", "钱七"],
        "age": [25, 30, 35, 40, 45],
        "salary": [10000, 15000, 20000, 25000, 30000],
        "department": ["研发部", "市场部", "财务部", "人事部", "行政部"]
    }
    _data_df = pd.DataFrame(data)

    owner_name = "叶磊老师"

    participants = {
        "user_001": "徐博",
        "user_002": "初京刚",
        "user_003": "于冰",
        "user_004": "珠兰",
        "user_005": "柴鸡",
        "user_006": "赵铜铁钢"
    }

    # 默认密钥文件路径
    default_key_file = "sm4_keys.json"

    demo_sm4_crypto(data_df=_data_df, owner_name=owner_name, participants=participants, key_file=default_key_file)
